import torch
import segmentation_models_pytorch as smp
import pytorch_lightning as pl

class SegModel(pl.LightningModule):
    def __init__(self, arch, encoder_name, in_channels, out_classes, t_max, mean, std, **kwargs):
        super().__init__()
        self.model = smp.create_model(
            arch, encoder_name=encoder_name, in_channels=in_channels, classes=out_classes, **kwargs
        )
        self.t_max = t_max
        self.number_of_classes = out_classes
        self.register_buffer("std", torch.tensor(std).view(1, in_channels, 1, 1))
        self.register_buffer("mean", torch.tensor(mean).view(1, in_channels, 1, 1))

    def forward(self, image):
        image = (image - self.mean) / self.std
        mask = self.model(image)
        return mask